/******************************************************************************
 * Encryptor class encrypts messages decryptable by the Decryptor class. It
 * uses OpenSSL enveloping to RSA-encrypt one-time symmetric session keys. The
 * session key is used to encrypt the payload. This work was done for Zero Leak
 * Architecture implementation for Dr. Hiroshi Fujinoki.
 *
 * Usage: Instantiated Encryptor must make the following calls in order:
 * keySrcFiles(...), encryptSessionKeys(), encrypt(...).
 *****************************************************************************/
#include "Encryptor.h"
#include "string.h"
#include "common.h"
#include "netinet/in.h"

Encryptor::Encryptor(const char** rsa_keyfiles, const char** expected_CN, int num_keys, const EVP_CIPHER* cipher) :
  type(cipher),
  npubk(0)
{
	ctx = new EVP_CIPHER_CTX;

	// random initialization vector ----
	iv = new unsigned char[EVP_MAX_IV_LENGTH];
	RAND_pseudo_bytes(iv, EVP_MAX_IV_LENGTH);

  keySrcFiles(rsa_keyfiles, num_keys);
  setupSessionKeys(expected_CN);
}

/******************************************************************************
 * KeySrcFiles: Locates rsa public key files and allocates memory for key-
 * 			related attributes.
 * Pre:		None.
 * Post:	Ready to call encryptSessionKeys().
 *****************************************************************************/
Encryptor& Encryptor::keySrcFiles(const char** rsa_keyfiles, int num_keys) {
	pubk_files = rsa_keyfiles;
	npubk = num_keys;

	ek = new unsigned char*[npubk];
	ekl = new int[npubk];
	pubk = new EVP_PKEY*[npubk];
	//rsa_keys = new RSA*[npubk];
	return *this;
}

/******************************************************************************
 * setupSessionKeys: Locates rsa public key files and allocates memory for
 * 			key-related attributes.
 * Pre: 	Must call keySrcFiles(...) first.
 * Post:	Ready to call encrypt(...).
 *****************************************************************************/
Encryptor& Encryptor::setupSessionKeys(const char** expected_CN)
{
  //BIO* rsa_bio;
  FILE* fp;
  X509* crt;

	for (int i = 0; i < npubk; i++)
  {
		// read in pem-encoded public key ----
		//rsa_bio = BIO_new_file(pubk_files[0], "r");
		//rsa_keys[i] = RSA_new();
		//PEM_read_bio_RSA_PUBKEY(rsa_bio, &rsa_keys[i], NULL, NULL);
		//BIO_free(rsa_bio);

    if (!(fp = fopen(pubk_files[i], "r")) ||
        !(crt = PEM_read_X509(fp, NULL, NULL, NULL)))
    {
        printf("Error reading Cert");
        exit(0);
    }

    //check cert
    if (!VerifyCertificate(crt, expected_CN[i]))
    {
       printf("Cert doesn't verify\n");
       exit(0);
    }

		// Assign PKEY
		//pubk[i] = EVP_PKEY_new();
		//EVP_PKEY_assign_RSA(pubk[i], rsa_keys[i]);
		pubk[i] = X509_get_pubkey(crt);

    X509_free(crt);
    fclose(fp);

		ek[i] = new unsigned char[EVP_PKEY_size(pubk[i])];
	}
	return *this;
}

/******************************************************************************
 * Encrypt: Performs encryption on message.
 * Pre: 	Must call encryptSessionKeys() first.
 * Post:	Encrypted message is stored in encrypt_buf, message length is
 * 			stored in decrypt_len, encrypt_buf length is stored in
 * 			encrypt_len, session keys are stored in ek, initialization vector
 * 			is stored in iv.
 *****************************************************************************/
unsigned char* Encryptor::encrypt(unsigned char* message, int message_len, int* encrypt_len) {

  int bytes_written;
	unsigned char* encrypt_buf = new unsigned char[message_len + AES_256_BLOCK_SIZE];

	EVP_SealInit(ctx, type, ek, ekl, iv, pubk, npubk);

	// EVP_SealUpdate changes overhead to # bytes in message ----
	EVP_SealUpdate(ctx, encrypt_buf, &bytes_written, message, message_len);
  *encrypt_len = bytes_written;

	// EVP_SealFinal puts the number of bytes written in the bytes_written int ----
	EVP_SealFinal(ctx, &encrypt_buf[*encrypt_len], &bytes_written);
	*encrypt_len += bytes_written;

  return (encrypt_buf);
}

unsigned char* Encryptor::MakeEncryptedPacket(unsigned char* msg, int msg_len, int* encrypt_len)
{
  int index = 0;
  int encrypted_data_len;
  uint32_t net_ekl, net_edl;

  unsigned char* encrypted_data = encrypt(msg, msg_len, &encrypted_data_len);

  *encrypt_len = encrypted_data_len + EVP_MAX_IV_LENGTH + ekl[0] + sizeof(uint32_t) + sizeof(uint32_t);

  unsigned char* encrypted_packet = new unsigned char[*encrypt_len];


  //iv
  memcpy(encrypted_packet, iv, EVP_MAX_IV_LENGTH);
  index += EVP_MAX_IV_LENGTH;

  //ekl
  net_ekl = ekl[0];
  net_ekl = htonl(net_ekl);
  memcpy(&encrypted_packet[index], &net_ekl, sizeof(uint32_t));
  index += sizeof(uint32_t);

  //ek
  memcpy(&encrypted_packet[index], ek[0], ekl[0]);
  index += ekl[0];

  //encrypted data size
  net_edl = encrypted_data_len;
  net_edl = htonl(net_edl);
  memcpy(&encrypted_packet[index], &net_edl, sizeof(uint32_t));
  index += sizeof(uint32_t);

  //Finally... encrypted data
  memcpy(&encrypted_packet[index], encrypted_data, encrypted_data_len);

  delete[] encrypted_data;

  return (encrypted_packet);
}

int Encryptor::getSessionKeyLength(int i) const {
	return ekl[i];
}

unsigned char* Encryptor::getSessionKey(int i) const {
	return ek[i];
}

unsigned char* Encryptor::getInitVector() const {
	return iv;
}

Encryptor::~Encryptor() {
	delete ctx;
	delete type;
	delete[] iv;
	for (int i = 0; i < npubk; i++) {
		EVP_PKEY_free(pubk[i]);
		delete ek[i];
		delete pubk_files[i];
	}
	// must come after loop ----
	delete[] ek;
	delete[] ekl;
	delete[] pubk;
	delete pubk_files;
	//delete[] rsa_keys;
}
